package com.fmync.ai.fmync.service;

import cn.hutool.core.io.FileUtil;
import com.fmync.ai.fmync.config.VectorStoreConfig;
import com.fmync.ai.fmync.utils.FileUploadUtil;
import dev.langchain4j.agent.tool.P;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.DocumentTransformer;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.data.segment.TextSegmentTransformer;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.*;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.util.List;

@Slf4j
@Service
@RequiredArgsConstructor
public class KnowledgeBaseService {
    private final OllamaStreamingChatModel streamingChatModel;
    private final OllamaEmbeddingModel ollamaEmbeddingModel;
    private final EmbeddingStore<TextSegment> embeddingStore;
    private final MessageWindowChatMemory chatMemory;
    private final VectorStoreConfig vectorStoreConfig;

    private final static List<String> OFFICETYPE = List.of("doc","docx","ppt","pptx","xls","xlsx");



    public void addDocument(MultipartFile multipartFile) {

        String path = null;
        String fileName = multipartFile.getOriginalFilename();
        String fileType = FileUtil.extName(fileName);
        try {
            path = FileUploadUtil.saveFile(multipartFile, "F:/home/langchain4j");
        } catch (IOException e) {
            throw new RuntimeException(e);
        }


        Document document = null;
        if("txt".equalsIgnoreCase(fileType)){
            document = FileSystemDocumentLoader.loadDocument(path, new TextDocumentParser());
        }else if("pdf".equalsIgnoreCase(fileType)){
            document = FileSystemDocumentLoader.loadDocument(path, new ApachePdfBoxDocumentParser());
        }else if(OFFICETYPE.contains(fileType.toLowerCase())){
            document = FileSystemDocumentLoader.loadDocument(path, new ApachePoiDocumentParser());
        }else{
            document = FileSystemDocumentLoader.loadDocument(path, new ApacheTikaDocumentParser());
        }


        EmbeddingStoreIngestor embeddingStoreIngestor =  EmbeddingStoreIngestor.builder()
                .documentSplitter(new DocumentByParagraphSplitter(500,50))
                .documentTransformer(new DocumentTransformer(){
                    @Override
                    public Document transform(Document document) {
                        document.metadata().put("fileName",fileName);
                        return document;
                    }
                })
                .embeddingModel(ollamaEmbeddingModel)
                .embeddingStore(embeddingStore)
                .textSegmentTransformer(new TextSegmentTransformer(){

                    @Override
                    public TextSegment transform(TextSegment textSegment) {
                        textSegment.metadata().put("fileName",fileName);
                        return textSegment;
                    }
                })
                .build();
        embeddingStoreIngestor.ingest(document);
    }

    public Flux<String> chatStream(String question) {
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(vectorStoreConfig.embeddingStore())
                .embeddingModel(ollamaEmbeddingModel)
                .maxResults(3)
                .minScore(0.75)
                .build();




        StreamingAssistant assistant = AiServices.builder(StreamingAssistant.class)
                .streamingChatLanguageModel(streamingChatModel)
//                .chatLanguageModel(chatModel)
                .contentRetriever(contentRetriever)
//                .chatMemory(chatMemory)
                .build();



        return  assistant.chatStream(UserMessage.from(question)).subscribeOn(Schedulers.boundedElastic()) ;

       /* return Flux.create(sink -> {
            assistant.chat(UserMessage.from(question))
                .onRetrieved((List<Content> sources) -> {
                    // 当相关源信息被检索到时触发此回调
                    // 可以在这里处理检索到的源信息
                    for (Content source : sources) {
                        System.out.println("应用文档名称："+source.metadata().get("fileName"));
                    }
                })
                .onPartialResponse(partialResponse -> {
                    // 当接收到部分响应时触发此回调
                    // partialResponse 包含了部分响应的内容
                    sink.next("1111");
                })
                .onCompleteResponse(completeResponse -> {
                    // 当接收到完整响应时触发此回调
                    // completeResponse 包含了完整的响应内容
                    sink.complete();
                })
                .onError(error -> {
                    // 当出现错误时触发此回调
                    // error 是具体的异常信息
                    sink.error(error);
                })
                .start();
        });*/
    }



} 